use super::*;
use riir::agent;
use riir::agent::tools::playground::{run_rustycli, VerifyResp, VerifyType};
use std::marker::PhantomData;
use tracing::info;

use serde::{Deserialize, Serialize};

pub const TRANS_RUST_PATH: &str = "./files/rust/trans.md";

#[derive(Serialize, Deserialize)]
pub struct RustTransInfo {
    pub rust_expert: String,
}

pub const RUST_TRANS: Prompt<RustTransInfo> = Prompt("RustTransInfo: [Rust Trans]", PhantomData);

#[derive(Serialize, Deserialize)]
pub struct TransPlanInfo {
    pub task: String,
}

pub const RUST_TRANS_PLAN: Prompt<TransPlanInfo> = Prompt(
    r#"
You are a programming language translation expert, focusing on converting existing C/C++ projects into Rust implementations. Your task is not merely to translate the code line by line, but to deeply understand the architecture and abstract concepts of the original project and reimplement these concepts in Rust. Please follow these guidelines for this task:

- **Please ensure the code is fully functional and doesn't omit any details**.
- **Understanding the Original Project's Architecture:** First, thoroughly analyze the overall architecture and module division of the C/C++ project. Identify key abstractions and data flows. This includes class hierarchies, interfaces, data models, and control flows.
- **Abstraction and Language Feature Mapping:** Determine how C/C++ language features can find their corresponding implementations in Rust. Pay special attention to how to handle pointers and references, memory management (e.g., using Rust’s ownership and borrowing rules instead of manual memory management), exception handling, etc., in C/C++.
- **Use of Stable Rust:** Ensure that your Rust implementation is based on the latest stable version of Rust, avoiding features that require nightly Rust. This is to ensure the stability and long-term maintainability of the code.
- **Performance and Safety:** During the conversion, consider Rust's performance optimizations and safety measures. Rust offers a powerful type system and ownership model to reduce memory safety issues. You should fully leverage these features to optimize the code.
- **Testing and Verification:** After conversion, write comprehensive unit and integration tests for the Rust code to ensure the new implementation's correctness and performance are consistent with or superior to the original C/C++ project.
- **Documentation and Code Style:** Write clear code comments and documentation according to Rust conventions.

Before starting the translation, plan your implementation strategy, considering how to proceed in phases and which parts may require special attention or restructuring. This task is not just about code translation but also a deep understanding and re-creation of the original project's architecture and design philosophy.

Please ensure that your reply strictly adheres to the following  JSON format:
```json
{{
	"plan": "<concise plan on how you will complete the task>",
	"code": "
	```rust
	<The Rust code after your audit, and the Rust code snippet without any escape characters Please!>
	```
	"
}}
```
"#,
    PhantomData,
);

#[derive(Debug)]
pub struct TransAgent {
    attributes: AgentBase,
    bug_errors: Option<String>,
    bug_count: u8,
    retry_count: u8,
    llm: Arc<Mutex<Box<dyn LLM + 'static>>>,
}

impl TransAgent {
    pub fn new(llm: Arc<Mutex<Box<dyn LLM + 'static>>>) -> Self {
        let attributes: AgentBase = AgentBase {
            objective: "Docs Rust Code".to_string(),
            position: "Docs Agent".to_string(),
            state: AgentState::Starting,
            memory: vec![],
        };

        Self {
            attributes,
            bug_errors: None,
            bug_count: 0,
            retry_count: 0,
            llm: llm.clone(),
        }
    }

    async fn call_initial_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        // Concatenate Instruction
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_TRANS_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "trans", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_improved_code(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_TRANS_PLAN.0),
            format!(
                "CODE: {} ",
                factsheet.code.as_ref().unwrap_or(&"".to_string())
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "trans", &ai_response).await;
        factsheet.code = Some(ai_response);
    }

    async fn call_fix_code_bugs(&mut self, factsheet: &mut FactSheet, output_path: PathBuf) {
        // let decode_code =
        //     decode_llm_output(factsheet.code.as_ref().unwrap_or(&"".to_string())).await;
        let msg_context: Vec<Message> = prompt::build_msg(
            format!("PLAN: {} ", RUST_FIXED_EXPERT.0),
            format!(
                "CODE: {:?} \n ERROR_BUGS: {:?} \n",
                factsheet.code.as_ref().unwrap_or(&"".to_string()),
                self.bug_errors
            ),
        );

        let ai_response: String =
            ai_task_request(msg_context, &self.attributes.position, self.llm.clone()).await;

        save_code(output_path, "trans", &ai_response).await;
        factsheet.code = Some(ai_response);
    }
}

#[async_trait]
impl AgentFns for TransAgent {
    fn get_attributes_from_agent(&self) -> &AgentBase {
        &self.attributes
    }

    async fn execute(
        &mut self,
        factsheet: &mut FactSheet,
        output_path: PathBuf,
    ) -> Result<(), Box<dyn std::error::Error>> {
        while self.attributes.state != AgentState::Done {
            match &self.attributes.state {
                AgentState::Starting => {
                    info!("Trans Agent Start ---");
                    self.call_initial_code(factsheet, output_path.clone()).await;
                    self.attributes.state = AgentState::Working;
                    continue;
                }

                AgentState::Working => {
                    info!("Trans Agent Working ---");
                    if self.bug_count == 0 {
                        self.call_improved_code(factsheet, output_path.clone()).await;
                    } else {
                        self.call_fix_code_bugs(factsheet, output_path.clone()).await;
                    }
                    self.attributes.state = AgentState::Debug;
                    continue;
                }

                AgentState::Debug => {
                    info!("Trans Agent Debug ---");

                    // --- 验证 LLM 响应代码
                    if let Some(e) = agent::capture_rust_code(&factsheet.code.clone().unwrap()) {
                        info!("EndANswer Chat: plain {:?}", e);

                        // do something async
                        let verft = run_rustycli(&e, VerifyType::BUILD).await;
                        info!("Verify Chat: {:?}", verft);
                        match verft {
                            VerifyResp::Success { stdout, stderr } => {
                                self.bug_count = 0;
                                info!("Code build successful...");
                                self.attributes.state = AgentState::Done;
                            }
                            VerifyResp::Fail(error_str) => {
                                // Update error stats
                                self.bug_count += 1;
                                self.retry_count += 1;
                                self.bug_errors = Some(error_str);
                                if self.retry_count > 10 {
                                    info!(
                                        "Code build failed but retry count is greater then 10... "
                                    );
                                    self.attributes.state = AgentState::Done;
                                } else {
                                    // Pass back for rework
                                    self.attributes.state = AgentState::Working;
                                    continue;
                                }
                            }
                        }
                    } else {
                        info!("Uncapture rust code!");
                        // Pass back for rework
                        self.attributes.state = AgentState::Working;
                        continue;
                    }
                }

                _ => {}
            }
        }
        Ok(())
    }
}
